import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.linalg import solve_continuous_are

class WheeledLegRobot:
    def __init__(self):
        # Robot parameters
        self.m_body = 2.0      # body mass kg
        self.m_leg = 0.5       # leg mass kg
        self.l_hip = 0.3       # hip joint length m
        self.l_knee = 0.3      # knee joint length m
        self.I_body = 0.1      # body moment of inertia
        self.g = 9.81          # gravity
        
        # State: [x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot]
        self.state_dim = 12
        self.control_dim = 6
        
    def dynamics(self, state, control_torques):
        """Robot dynamics equations"""
        x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot = state
        tau1, tau2, tau3, tau4, tau5, tau6 = control_torques
        
        # Simplified dynamics model
        M = np.diag([self.m_body, self.m_body, self.m_body, 
                     self.I_body, self.I_body, self.I_body])
        
        # Coriolis and gravity terms
        C = np.array([0, 0, self.m_body * self.g, 0, 0, 0])
        
        # Control input mapping
        B = np.eye(6)
        
        # Acceleration calculation
        accelerations = np.linalg.solve(M, B @ control_torques - C)
        
        # State derivatives
        state_dot = np.zeros(12)
        state_dot[0:6] = state[6:12]  # position derivatives = velocities
        state_dot[6:12] = accelerations  # velocity derivatives = accelerations
        
        return state_dot

class PIDController:
    def __init__(self, kp, ki, kd, dt):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.dt = dt
        self.integral = np.zeros(6)
        self.previous_error = np.zeros(6)
    
    def compute_control(self, desired_state, current_state):
        error = desired_state[0:6] - current_state[0:6]
        error_derivative = (error - self.previous_error) / self.dt
        
        self.integral += error * self.dt
        
        control = (self.kp * error + 
                  self.ki * self.integral + 
                  self.kd * error_derivative)
        
        self.previous_error = error.copy()
        return control

class SMCController:
    def __init__(self, lambda_val, k, phi, dt):
        self.lambda_val = lambda_val
        self.k = k
        self.phi = phi
        self.dt = dt
        self.previous_error = np.zeros(6)
    
    def compute_control(self, desired_state, current_state):
        error = desired_state[0:6] - current_state[0:6]
        error_derivative = (error - self.previous_error) / self.dt
        
        # Sliding surface
        s = error_derivative + self.lambda_val * error
        
        # Control law (using saturation function to reduce chattering)
        sat_s = np.tanh(s * 10)  # saturation function instead of sign function
        control = self.k * sat_s + self.phi * s
        
        self.previous_error = error.copy()
        return control

class LQRController:
    def __init__(self, A, B, Q, R):
        # Use scipy to solve continuous-time algebraic Riccati equation
        P = solve_continuous_are(A, B, Q, R)
        self.K = np.linalg.solve(R, B.T @ P)
    
    def compute_control(self, desired_state, current_state):
        error = current_state - desired_state
        return -self.K @ error

def simulate_controller(controller, robot, desired_trajectory, t_span, dt):
    """Simulate specific controller"""
    t_eval = np.arange(t_span[0], t_span[1], dt)
    n_steps = len(t_eval)
    
    # Initialize state records
    states = np.zeros((n_steps, robot.state_dim))
    controls = np.zeros((n_steps, robot.control_dim))
    errors = np.zeros((n_steps, 6))
    
    # Initial state
    state = np.array([0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0])  # initial height 0.5m
    
    for i, t in enumerate(t_eval):
        # Get desired state
        desired_state = desired_trajectory(t)
        
        # Calculate control input
        control = controller.compute_control(desired_state, state)
        
        # Limit control input
        control = np.clip(control, -20, 20)
        
        # Record state and control
        states[i] = state
        controls[i] = control
        errors[i] = desired_state[0:6] - state[0:6]
        
        # Integrate dynamics (using Euler method for simplification)
        state_dot = robot.dynamics(state, control)
        state = state + state_dot * dt
    
    return t_eval, states, controls, errors

def desired_circular_trajectory(t):
    """Circular trajectory"""
    radius = 1.0
    omega = 0.5
    
    x = radius * np.cos(omega * t)
    y = radius * np.sin(omega * t)
    z = 0.5  # maintain fixed height
    theta, phi, psi = 0, 0, 0  # maintain horizontal attitude
    x_dot = -radius * omega * np.sin(omega * t)
    y_dot = radius * omega * np.cos(omega * t)
    z_dot = 0
    theta_dot, phi_dot, psi_dot = 0, 0, 0
    
    return np.array([x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot])

def plot_3d_trajectory(states_pid, states_smc, states_lqr):
    """Plot 3D trajectory comparison"""
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # Plot trajectories
    ax.plot(states_pid[:, 0], states_pid[:, 1], states_pid[:, 2], 
            'r-', label='PID', linewidth=2, alpha=0.8)
    ax.plot(states_smc[:, 0], states_smc[:, 1], states_smc[:, 2], 
            'g-', label='SMC', linewidth=2, alpha=0.8)
    ax.plot(states_lqr[:, 0], states_lqr[:, 1], states_lqr[:, 2], 
            'b-', label='LQR', linewidth=2, alpha=0.8)
    
    # Plot desired trajectory
    t_ref = np.linspace(0, 10, 100)
    x_ref = np.cos(0.5 * t_ref)
    y_ref = np.sin(0.5 * t_ref)
    z_ref = 0.5 * np.ones_like(t_ref)
    ax.plot(x_ref, y_ref, z_ref, 'k--', label='Desired', linewidth=2)
    
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_zlabel('Z (m)')
    ax.set_title('3D Trajectory Tracking Comparison')
    ax.legend()
    ax.grid(True)
    
    # Set equal axis scale
    max_range = 1.5
    ax.set_xlim([-max_range, max_range])
    ax.set_ylim([-max_range, max_range])
    ax.set_zlim([0, max_range])

# Main simulation
def main():
    # Initialize robot and controllers
    robot = WheeledLegRobot()
    dt = 0.01
    t_span = (0, 10)
    
    # PID Controller
    pid = PIDController(
        kp=np.array([15, 15, 20, 10, 10, 10]),
        ki=np.array([0.1, 0.1, 0.5, 0.1, 0.1, 0.1]),
        kd=np.array([5, 5, 8, 3, 3, 3]),
        dt=dt
    )
    
    # SMC Controller
    smc = SMCController(
        lambda_val=2.0,
        k=np.array([8, 8, 10, 5, 5, 5]),
        phi=np.array([2, 2, 3, 1, 1, 1]),
        dt=dt
    )
    
    # LQR Controller (linearized model)
    A = np.zeros((12, 12))
    A[0:6, 6:12] = np.eye(6)
    A[8, 8] = -0.1  # simple damping term
    
    B = np.zeros((12, 6))
    B[6:12, 0:6] = np.eye(6) / robot.m_body
    B[9:12, 3:6] = np.eye(3) / robot.I_body
    
    Q = np.diag([10, 10, 20, 5, 5, 5, 1, 1, 2, 0.5, 0.5, 0.5])
    R = np.diag([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
    
    lqr = LQRController(A, B, Q, R)
    
    # Run simulations
    print("Running PID controller simulation...")
    t_pid, states_pid, controls_pid, errors_pid = simulate_controller(
        pid, robot, desired_circular_trajectory, t_span, dt)
    
    print("Running SMC controller simulation...")
    t_smc, states_smc, controls_smc, errors_smc = simulate_controller(
        smc, robot, desired_circular_trajectory, t_span, dt)
    
    print("Running LQR controller simulation...")
    t_lqr, states_lqr, controls_lqr, errors_lqr = simulate_controller(
        lqr, robot, desired_circular_trajectory, t_span, dt)
    
    # Plot 3D trajectory
    plot_3d_trajectory(states_pid, states_smc, states_lqr)
    
    # Plot other comparison charts
    plot_comparison(t_pid, t_smc, t_lqr, 
                   states_pid, states_smc, states_lqr,
                   errors_pid, errors_smc, errors_lqr,
                   controls_pid, controls_smc, controls_lqr)

def plot_comparison(t_pid, t_smc, t_lqr, states_pid, states_smc, states_lqr, 
                   errors_pid, errors_smc, errors_lqr, controls_pid, controls_smc, controls_lqr):
    """Plot comparison results of three controllers"""
    
    fig, axes = plt.subplots(3, 3, figsize=(15, 12))
    
    # Position tracking comparison
    axes[0,0].plot(t_pid, states_pid[:, 0], 'r-', label='PID', linewidth=2)
    axes[0,0].plot(t_smc, states_smc[:, 0], 'g-', label='SMC', linewidth=2)
    axes[0,0].plot(t_lqr, states_lqr[:, 0], 'b-', label='LQR', linewidth=2)
    
    # Plot desired trajectory
    t_ref = np.linspace(0, 10, 100)
    x_ref = np.cos(0.5 * t_ref)
    axes[0,0].plot(t_ref, x_ref, 'k--', label='Desired', linewidth=1)
    
    axes[0,0].set_ylabel('X Position (m)')
    axes[0,0].set_title('X-Axis Position Tracking')
    axes[0,0].legend()
    axes[0,0].grid(True)
    
    axes[0,1].plot(t_pid, states_pid[:, 1], 'r-', label='PID', linewidth=2)
    axes[0,1].plot(t_smc, states_smc[:, 1], 'g-', label='SMC', linewidth=2)
    axes[0,1].plot(t_lqr, states_lqr[:, 1], 'b-', label='LQR', linewidth=2)
    
    y_ref = np.sin(0.5 * t_ref)
    axes[0,1].plot(t_ref, y_ref, 'k--', label='Desired', linewidth=1)
    
    axes[0,1].set_ylabel('Y Position (m)')
    axes[0,1].set_title('Y-Axis Position Tracking')
    axes[0,1].legend()
    axes[0,1].grid(True)
    
    axes[0,2].plot(t_pid, states_pid[:, 2], 'r-', label='PID', linewidth=2)
    axes[0,2].plot(t_smc, states_smc[:, 2], 'g-', label='SMC', linewidth=2)
    axes[0,2].plot(t_lqr, states_lqr[:, 2], 'b-', label='LQR', linewidth=2)
    axes[0,2].set_ylabel('Z Position (m)')
    axes[0,2].set_title('Z-Axis Height Tracking')
    axes[0,2].legend()
    axes[0,2].grid(True)
    
    # Tracking error comparison
    error_norm_pid = np.linalg.norm(errors_pid[:, 0:3], axis=1)
    error_norm_smc = np.linalg.norm(errors_smc[:, 0:3], axis=1)
    error_norm_lqr = np.linalg.norm(errors_lqr[:, 0:3], axis=1)
    
    axes[1,0].plot(t_pid, error_norm_pid, 'r-', label='PID', linewidth=2)
    axes[1,0].plot(t_smc, error_norm_smc, 'g-', label='SMC', linewidth=2)
    axes[1,0].plot(t_lqr, error_norm_lqr, 'b-', label='LQR', linewidth=2)
    axes[1,0].set_ylabel('Position Error Norm (m)')
    axes[1,0].set_title('Position Error Norm')
    axes[1,0].legend()
    axes[1,0].grid(True)
    
    # Control input comparison
    control_norm_pid = np.linalg.norm(controls_pid, axis=1)
    control_norm_smc = np.linalg.norm(controls_smc, axis=1)
    control_norm_lqr = np.linalg.norm(controls_lqr, axis=1)
    
    axes[1,1].plot(t_pid, control_norm_pid, 'r-', label='PID', linewidth=2)
    axes[1,1].plot(t_smc, control_norm_smc, 'g-', label='SMC', linewidth=2)
    axes[1,1].plot(t_lqr, control_norm_lqr, 'b-', label='LQR', linewidth=2)
    axes[1,1].set_ylabel('Control Effort Norm')
    axes[1,1].set_title('Control Input Norm')
    axes[1,1].legend()
    axes[1,1].grid(True)
    
    # Trajectory tracking performance
    axes[1,2].plot(states_pid[:, 0], states_pid[:, 1], 'r-', label='PID', linewidth=2, alpha=0.7)
    axes[1,2].plot(states_smc[:, 0], states_smc[:, 1], 'g-', label='SMC', linewidth=2, alpha=0.7)
    axes[1,2].plot(states_lqr[:, 0], states_lqr[:, 1], 'b-', label='LQR', linewidth=2, alpha=0.7)
    
    # Plot desired trajectory
    axes[1,2].plot(x_ref, y_ref, 'k--', label='Desired', linewidth=2)
    
    axes[1,2].set_xlabel('X (m)')
    axes[1,2].set_ylabel('Y (m)')
    axes[1,2].set_title('XY Plane Trajectory Tracking')
    axes[1,2].legend()
    axes[1,2].grid(True)
    axes[1,2].axis('equal')
    
    # Performance metrics comparison
    metrics = {
        'RMSE Position': [
            np.sqrt(np.mean(error_norm_pid**2)),
            np.sqrt(np.mean(error_norm_smc**2)),
            np.sqrt(np.mean(error_norm_lqr**2))
        ],
        'Max Control Effort': [
            np.max(control_norm_pid),
            np.max(control_norm_smc),
            np.max(control_norm_lqr)
        ],
        'Mean Control Effort': [
            np.mean(control_norm_pid),
            np.mean(control_norm_smc),
            np.mean(control_norm_lqr)
        ]
    }
    
    controllers = ['PID', 'SMC', 'LQR']
    x_pos = np.arange(len(controllers))
    width = 0.25
    
    for i, (metric_name, values) in enumerate(metrics.items()):
        bars = axes[2,i].bar(x_pos, values, width, 
                           color=['red', 'green', 'blue'], alpha=0.7)
        axes[2,i].set_xticks(x_pos)
        axes[2,i].set_xticklabels(controllers)
        axes[2,i].set_title(metric_name)
        axes[2,i].grid(True, alpha=0.3)
        
        # Add value labels
        for j, v in enumerate(values):
            axes[2,i].text(j, v + 0.1, f'{v:.3f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.show()
    
    # Output performance comparison
    print("\n=== Performance Metrics Comparison ===")
    print(f"{'Metric':<20} {'PID':<10} {'SMC':<10} {'LQR':<10}")
    print("-" * 50)
    for metric, values in metrics.items():
        print(f"{metric:<20} {values[0]:<10.3f} {values[1]:<10.3f} {values[2]:<10.3f}")

if __name__ == "__main__":
    main()